#! /usr/bin/env python
# coding=utf-8
# 建立一个用于上传的文件

import json
import argparse
import sys
from argparse import RawTextHelpFormatter

parser = argparse.ArgumentParser(
    description='''
    整合 emmapper gff3 gene2swiss table 和 gene exp 文件 生成一个用于导入我们自己的数据库的文件

    用法:
    make_one_file_for_upload.py -g Cves_OGS_v3.7.gff3 -s out_gene2swiss_table -e gene_exp.tsv -m emapper_out.emapper.annotations -o all_genes.tsv


    ''',formatter_class=RawTextHelpFormatter)


parser.add_argument('-g',
                help='必须给定，gff3 文件')
parser.add_argument('-s',
                help='可选， 由match_swiss_prot_annotation.py 生成的 swiss 表格')
parser.add_argument('-e',
                help='可选， 由read_rpkm_from_all_sample.py  生成的 表达量表格')
parser.add_argument('-m',
                help='可选， emapper 注释文件')

parser.add_argument('-o',
                help='输出的tsv文件')


args = parser.parse_args()

if not args.g or not args.o:
    parser.print_help()
    sys.exit()



if not args.s:
	infile_swiss = None
else:
	infile_swiss = args.s

if not args.e:
	infile_gene_exp = None
else:
	infile_gene_exp = args.e


if not args.e:
	infile_emapper = None
else:
	infile_emapper = args.m



infile_gff = args.g

outfile = args.o


go_dic = {}
kegg_dic = {}
swiss_dic = {}
gene_exp_dic = {}
gene_gff_dic={}


if infile_emapper != None:

	with open(infile_emapper) as fila:
		
		for i in fila:
			k = i.split('#')[0].split('\t')
			if len(k)>2:
				gene_name = k[0].split('.mRNA')[0]
				if k[9] == '-':
					k[9] = ''
				if k[11] =='-':
					k[11] = ''

				go_lista = [a for a in k[9].split(',') if a != '']
				#print(k[11])
				
				kegg_lista = [a.split('ko:')[1] for a in k[11].split(',') if len(a.split('ko:'))>1]
				#print(kegg_lista)
				go_dic[gene_name] = go_lista
				kegg_dic[gene_name] = kegg_lista 

if infile_gene_exp != None:
	with open(infile_gene_exp) as fila:
		for i in fila:
			k = i.strip().split('\t')
			if len(k)>1:
				gene_exp_dic[k[0]] = k[1] 

# 读取

with open(infile_gff) as fila:
	for i in fila:
		k = i.strip().split('\t')
		#print(k)
		if len(k)>5 and k[2] == 'gene':
			name = k[8].split('Name=')[1].split(';')[0]
			gene_gff_dic[name] = [k[0],k[6],k[3],k[4]]
		
			#print([k[0],k[6],k[3],k[4]])


if infile_swiss != None:
	with open(infile_swiss) as fita:
		for i in fita:
			k = i.strip().split('\t')
			if len(k)>1:
				gene_name = k[0]
				swiss_name = k[1]
				des = k[3]
				sym = k[2]
				swiss_dic[gene_name] = [swiss_name,sym,des]


with open(outfile,'w') as fila:
	fila.write('gene_id\tgene_acc\tgene_name\ttype\tgene_symbol\tgene_description\tswissprot\tseq_name\tdirection\tstart\tend\tKEGG\tGO\tgene_expression\n')

	for gene_name in gene_gff_dic:
		swiss_target = ['','','']
		if gene_name in swiss_dic:
			swiss_target = swiss_dic[gene_name]
		
		go_target = []
		if gene_name in go_dic:
			go_target = go_dic[gene_name]

		kegg_target = []
		if gene_name in kegg_dic:
			kegg_target = kegg_dic[gene_name]

		gene_exp_target = ''
		if gene_name in gene_exp_dic:
			gene_exp_target = gene_exp_dic[gene_name]


		fila.write('\t'+gene_name+'\t'+gene_name+'\tprotein coding gene\t'+swiss_target[1]+'\t'+swiss_target[2]+'\t'+swiss_target[0]+'\t'+
			gene_gff_dic[gene_name][0]+'\t'+gene_gff_dic[gene_name][1]+'\t'+ gene_gff_dic[gene_name][2]+'\t'+gene_gff_dic[gene_name][3]+'\t'
			+json.dumps(kegg_target)+'\t'+json.dumps(go_target)+'\t'+gene_exp_target+'\n')





